import os
import pandas as pd
from tqdm import tqdm
import shutil
import re
import argparse
from frontend.zh_frontend import Frontend
# data_root = "./data_aishell3"
# work_path = "./mfa_train_data"


def refresh_dir(path, overwrite=False, parent=False):
    if parent:

        path = os.path.dirname(path)

    if not os.path.exists(path):

        os.makedirs(path)

    else:
        if overwrite:
            shutil.rmtree(path)
            os.makedirs(path)
        



def gen_audio_file_text_dic(table_file: str):
    if table_file.endswith("csv"):
        df = pd.read_csv(table_file)
    elif table_file.endswith("xlsx") or table_file.endswith("xls"):

        df = pd.read_excel(table_file)

    else:
        raise Exception(table_file+" not supported")

    audio_file_text_dic = {}

    for audio_file, text in zip(df["audio_file"], df["text"]):

        audio_file_text_dic[audio_file] = text

    return audio_file_text_dic


def prepare_mfa_train_files(
    data_root,
    linked_wav_root,
    mfa_lab_path,
    wav_text_table,
    lexicon_path,
    data_info_table = None
):
    if data_info_table:

        audio_name_list = []
        audio_path_list = []
        speaker_list = []
        text_list = []
        lab_list = []
        py_list = []
    frontend = Frontend()

    audio_dic = gen_audio_file_text_dic(wav_text_table)
    refresh_dir(linked_wav_root,overwrite=True)
    refresh_dir(mfa_lab_path,overwrite=True)
    lexicon_set = set()
    for top, _, nondirs in tqdm(list(os.walk(data_root)),desc="gen lab file and wav links"):
        for file in nondirs:
            if file.endswith(".wav"):
                text = audio_dic.get(file)
                if text is None:
                    continue
                phonemes = frontend.get_phonemes(text)[0]
                for p_ in phonemes:

                    lexicon_set.add(p_)
                lab = " ".join(phonemes)
                

                source_wav = os.path.join(top, file)
                linked_path = re.sub(r"^"+data_root,
                                     linked_wav_root, source_wav)

                lab_path = re.sub(r"^"+data_root, mfa_lab_path, source_wav)
                lab_path = re.sub(r".wav$", ".lab", lab_path)
                refresh_dir(linked_path, parent=True)
                os.symlink(source_wav, linked_path)

                refresh_dir(lab_path, parent=True)
                with open(lab_path, mode="w") as wf:
                    wf.write(lab + '\n')

                audio_name = file.replace(".wav", "")
                text = audio_dic.get(file)
                if text is None:
                    continue
                if data_info_table:

                    py_ = " ".join(frontend.text2pinyin(text)[0])
                    audio_path_list.append(os.path.join(top, file))
                    audio_name_list.append(audio_name)
                    speaker_list.append(top.split("/")[-1])
                    text_list.append(text)
                    lab_list.append(lab)
                    py_list.append(py_)
    if data_info_table:
        df = pd.DataFrame({
            "audio_name": audio_name_list,
            "audio_path": audio_path_list,
            "speaker": speaker_list,
            "text": text_list,
            "lab": lab_list,
            "py": py_list
        })

        df.to_csv(data_info_table, index=False)

    with open(lexicon_path , mode="w") as f:
        lines = [" ".join([item,item])+"\n" for item in lexicon_set]

        f.writelines(lines)



if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="prepare MFA corpus_directory lab and pronunciation dictionary for Chinese pinyin to phoneme for MFA inclued light tone and er hua yin")
    parser.add_argument("data_root", type=str, help="wav data root ")
    parser.add_argument("work_path", type=str,
                        help="path to store prepare files ")
    parser.add_argument("wav_text_table", type=str,
                        help="csv file or excel file ,  wav_file as column 'audio_file' without parent dir , text content as column 'text' ")
    parser.add_argument("linked_wav_root", type=str,help="path to slink wav under work_path")
    parser.add_argument("mfa_lab_path", type=str,help="path to gen phoneme .lab file ")
    parser.add_argument("lexicon_path", type=str,help="path to gen lexicon file ")
    

    args = parser.parse_args()
    
    DATA_INFO_TABLE = "data_info.csv"
    LINKED_WAV_ROOT = args.linked_wav_root
    MFA_LAB_PATH = args.mfa_lab_path
    lexicon_path =  args.lexicon_path
    
    LINKED_WAV_ROOT = os.path.join(args.work_path, LINKED_WAV_ROOT)
    MFA_LAB_PATH = os.path.join(args.work_path, MFA_LAB_PATH)
    DATA_INFO_TABLE = os.path.join(args.work_path, DATA_INFO_TABLE)
    lexicon_path = os.path.join(args.work_path, lexicon_path)
    data_root = os.path.abspath(args.data_root)
    LINKED_WAV_ROOT = os.path.abspath(LINKED_WAV_ROOT)
    MFA_LAB_PATH = os.path.abspath(MFA_LAB_PATH)
    lexicon_path = os.path.abspath(lexicon_path)

    prepare_mfa_train_files(data_root, LINKED_WAV_ROOT,
                            MFA_LAB_PATH, args.wav_text_table,lexicon_path, DATA_INFO_TABLE)
    print("Done!")
